package com.atwulidun.huffmancode.zipfile;

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.*;

public class HuffmanCodeDemo {
    public static void main(String[] args) {
        // 一、对一个文件进行压缩
        HuffmanCode.zipFile("E:\\huffmanCodeUnzippedFile\\1608706642.png", "E:\\huffmanCodeZippedFile\\zippedImage.png");
        // 二、对一个文件进行解压
        HuffmanCode.unzipFile("E:\\huffmanCodeZippedFile\\zippedImage.png", "E:\\huffmanCodeUnzippedFile\\unzippedImage.png");
    }
}

class HuffmanCode implements Serializable{
    private Map<Byte, String> map = new HashMap<>();
    private StringBuilder stringBuilder = new StringBuilder();
    private int leng;

    public Map<Byte, String> getMap() {
        return map;
    }

    public void setMap(Map<Byte, String> map) {
        this.map = map;
    }

    public StringBuilder getStringBuilder() {
        return stringBuilder;
    }

    public void setStringBuilder(StringBuilder stringBuilder) {
        this.stringBuilder = stringBuilder;
    }

    public int getLeng() {
        return leng;
    }

    public void setLeng(int leng) {
        this.leng = leng;
    }

    private List<Node> getListNode(byte[] bytes) {
        Map<Byte, Integer> hashMap = new HashMap<>();
        for (byte aByte : bytes) {
            Integer counts = hashMap.get(aByte);
            if (counts == null) {
                hashMap.put(aByte, 1);
            } else {
                hashMap.put(aByte, counts + 1);
            }
        }
        List<Node> list = new ArrayList<>();
        for (Map.Entry<Byte, Integer> byteIntegerEntry : hashMap.entrySet()) {
            list.add(new Node(byteIntegerEntry.getKey(), byteIntegerEntry.getValue()));
        }
        return list;
    }

    private Node getHuffmanTree(List<Node> list) {
        while (list.size() > 1) {
            Collections.sort(list);
            Node left = list.get(0);
            Node right = list.get(1);
            Node parent = new Node(null, left.getWeight() + right.getWeight());
            parent.setLeft(left);
            parent.setRight(right);
            list.remove(left);
            list.remove(right);
            list.add(parent);
        }
        return list.get(0);
    }

    private void getHuffmanCode(Node node) {
        // 这里写成node.getData() != 0会报空指针异常
        // 原因：Integer a = new Integer(66);// 手动装箱
        //      int i = a.intValue();// 手动拆箱
        // 比较的时候会自动拆箱，则node.getData() != 0的实际运行过程为：
        // Integer data = node.getData();
        // data == null;
        // int dataInt = data.intValue();// 引发空指针异常
        if (node != null && node.getData() != null) {
            map.put(node.getData(), stringBuilder.toString());
            return;
        }
        stringBuilder.append("0");
        getHuffmanCode(node.getLeft());
        stringBuilder.deleteCharAt(stringBuilder.length() - 1);

        stringBuilder.append("1");
        getHuffmanCode(node.getRight());
        stringBuilder.deleteCharAt(stringBuilder.length() - 1);
    }

    private byte[] huffmanZip(byte[] sourceBytes) {
        getHuffmanCode(getHuffmanTree(getListNode(sourceBytes)));
        StringBuilder stringBuilder = new StringBuilder();
        for (byte aByte : sourceBytes) {
            stringBuilder.append(map.get(aByte));
        }
        System.out.println("压缩时赫夫曼字符串的长度为：" + stringBuilder.length());
        int len = (stringBuilder.length() + 7) / 8;
        byte[] huffmanBytes = new byte[len];
        int index = 0;
        for (int i = 0; i < stringBuilder.length(); i += 8) {
            String huffmanStr = null;
            if (i + 8 > stringBuilder.length()) {
                huffmanStr = stringBuilder.substring(i);
            } else {
                huffmanStr = stringBuilder.substring(i, i + 8);
                if (i == stringBuilder.length() - 8) {
                    System.out.println(huffmanStr);
                }
            }
            huffmanBytes[index++] = (byte) Integer.parseInt(huffmanStr, 2);
        }
        leng = stringBuilder.length() % 8;
        // 要特别特别注意：如果赫夫曼字符串的长度能够被8整除，且最后八个字符为0开头，那么要把leng置为8
        // 不然会出现bug：01111111，leng == 0，那么在恢复赫夫曼字符串的时候就为1111111，少了前面的0
        if (leng == 0) {
            leng = 8;
        }
        return huffmanBytes;
    }

    private String byteToBitString(boolean flag, byte b) {
        int temp = b;
        if (flag) {
            temp |= 256;
        }
        String str = Integer.toBinaryString(temp);
        if (flag) {
            return str.substring(str.length() - 8);
        } else {
            if (b < 0) {
                return str.substring(str.length() - 8);
            } else {
                while (str.length() < leng) {
                    str = "0" + str;
                }
                return str;
            }
        }
    }

    private byte[] huffmanUnzip(byte[] huffmanBytes) {
        StringBuilder huffmanStr = new StringBuilder();
        for (int i = 0; i < huffmanBytes.length; i++) {
            boolean flag = (i == huffmanBytes.length - 1);
            byte b = huffmanBytes[i];
            huffmanStr.append(byteToBitString(!flag, b));
        }
        System.out.println("解压时赫夫曼字符串的长度为：" + huffmanStr.length());
        Map<String, Byte> hashMap = new HashMap<>();
        for (Map.Entry<Byte, String> byteStringEntry : map.entrySet()) {
            hashMap.put(byteStringEntry.getValue(), byteStringEntry.getKey());
        }
        List<Byte> list = new ArrayList<>();
        int i = 0;
        while (i < huffmanStr.length()) {
            int count = 1;
            Byte b = null;
            while (true) {
                String key = huffmanStr.substring(i, i + count);
                b = hashMap.get(key);
                if (b == null) {
                    count++;
                } else {
                    list.add(b);
                    break;
                }
            }
            i += count;
        }
        byte[] sourceBytes = new byte[list.size()];
        for (int i1 = 0; i1 < list.size(); i1++) {
            sourceBytes[i1] = list.get(i1);
        }
        return sourceBytes;
    }

    // 以下两个为压缩文件和解压缩文件的方法
    // 方法一：读取指定位置未压缩文件，压缩后写入到指定位置(静态工具方法)
    public static void zipFile(String source, String destination) {
        // 首先要创建一个HuffmanCode对象，并用其来调用huffmanZip方法，这样赫夫曼编码等数据便保存在了该对象中
        HuffmanCode huffmanCode = new HuffmanCode();
        try (InputStream is = new FileInputStream(source);// 读取文件
             OutputStream os = new FileOutputStream(destination);// 写入对象
             ObjectOutputStream oos = new ObjectOutputStream(os)) {// 写入对象
            // 读取文件字节数组
            System.out.println("读取到的字节数组的长度为：" + is.available());
            byte[] sourceBytes = new byte[is.available()];
            is.read(sourceBytes);// 这么重要的一句话居然漏掉了！！！！！！！！！！！！！！
            // 将文件字节数组转换成赫夫曼字节数组
            byte[] huffmanBytes = huffmanCode.huffmanZip(sourceBytes);
            System.out.println("压缩时赫夫曼数组的长度为：" + huffmanBytes.length);
            // 将赫夫曼字节数组以及HuffmanCode对象利用对象流的方式依次写入到指定位置
            // 因为HuffmanCode对象保存的数据在解压缩文件时需要用到，故也要一并写过去
            oos.writeObject(huffmanBytes);
            oos.writeObject(huffmanCode);
            oos.flush();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    // 方法二：读取指定位置已压缩文件，解压后写入到指定位置(静态工具方法)
    public static void unzipFile(String source, String destination) {
        try (InputStream is = new FileInputStream(source);// 读取对象
             ObjectInputStream ois = new ObjectInputStream(is);// 读取对象
             OutputStream os = new FileOutputStream(destination)) {// 写入文件
            // 使用对象流的方式可以分别读取之前用对象流写入的对象
            // 先读取huffmanBytes
            byte[] huffmanBytes = (byte[]) ois.readObject();
            System.out.println("解压时获得的赫夫曼数组的长度为：" + huffmanBytes.length);
            // 再读取HuffmanCode对象
            HuffmanCode huffmanCode = (HuffmanCode) ois.readObject();
            System.out.println(huffmanCode.getLeng());
            // 利用huffmanCode来调用huffmanUnzip方法
            byte[] sourceBytes = huffmanCode.huffmanUnzip(huffmanBytes);
            // 将解压得到的sourceBytes写入到指定的位置
            os.write(sourceBytes);
        } catch (IOException | ClassNotFoundException e) {
            e.printStackTrace();
        }
    }

}

class Node implements Comparable<Node> {
    private Byte data;
    private Integer weight;
    private Node left;
    private Node right;

    public Node(Byte data, Integer weight) {
        this.data = data;
        this.weight = weight;
    }

    public Byte getData() {
        return data;
    }

    public void setData(Byte data) {
        this.data = data;
    }

    public Integer getWeight() {
        return weight;
    }

    public void setWeight(Integer weight) {
        this.weight = weight;
    }

    public Node getLeft() {
        return left;
    }

    public void setLeft(Node left) {
        this.left = left;
    }

    public Node getRight() {
        return right;
    }

    public void setRight(Node right) {
        this.right = right;
    }

    @Override
    public String toString() {
        return "Node{" +
                "data=" + data +
                ", weight=" + weight +
                '}';
    }

    @Override
    public int compareTo(Node o) {
        return weight - o.weight;
    }
}
